from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import xml.etree.ElementTree as ET
import os
import pickle
import numpy as np
import pdb

def conf_eval(detpath,
             annopath,
             imagesetfile,
             classname,
             cachedir,
             ovthresh=0.5):
  """rec, prec, ap = voc_eval(detpath,
                              annopath,
                              imagesetfile,
                              classname,
                              [ovthresh],
                              [use_07_metric])

  Top level function that does the PASCAL VOC evaluation.

  detpath: Path to detections
      detpath.format(classname) should produce the detection results file.
  annopath: Path to annotations
      annopath.format(imagename) should be the xml annotations file.
  imagesetfile: Text file containing the list of images, one image per line.
  classname: Category name (duh)
  cachedir: Directory for caching the annotations
  [ovthresh]: Overlap threshold (default = 0.5)
  """
  # assumes detections are in detpath.format(classname)
  # assumes annotations are in annopath.format(imagename)
  # assumes imagesetfile is a text file with each line an image name
  # cachedir caches the annotations in a pickle file

  # first load gt
  if not os.path.isdir(cachedir):
    os.mkdir(cachedir)
  cachefile = os.path.join(cachedir, 'annots.pkl')

  # read list of images
  with open(imagesetfile, 'r') as f:
    lines = f.readlines()
  imagenames = [x.strip() for x in lines]

  if not os.path.isfile(cachefile):
    # load annotations
    recs = {}
    for i, imagename in enumerate(imagenames):
      recs[imagename] = parse_rec(annopath.format(imagename))
      if i % 100 == 0:
        print('Reading annotation for {:d}/{:d}'.format(
          i + 1, len(imagenames)))
    # save
    print('Saving cached annotations to {:s}'.format(cachefile))
    with open(cachefile, 'wb') as f:
      pickle.dump(recs, f)
  else:
    # load
    with open(cachefile, 'rb') as f:
      try:
        recs = pickle.load(f)
      except:
        recs = pickle.load(f, encoding='bytes')

  # extract gt objects for this class
  class_recs = {}
  npos = 0
  for imagename in imagenames:
    R = [obj for obj in recs[imagename] if obj['name'] == classname]
    bbox = np.array([x['bbox'] for x in R])
    difficult = np.array([x['difficult'] for x in R]).astype(np.bool)
    det = [False] * len(R)
    npos = npos + sum(~difficult)
    class_recs[imagename] = {'bbox': bbox,
                             'difficult': difficult,
                             'det': det}

  # read dets
  detfile = detpath.format(classname)
  with open(detfile, 'r') as f:
    lines = f.readlines()

  if len(lines) == 0:
        # No detection examples
        return 0, 0, 0, 0, npos

  splitlines = [x.strip().split(' ') for x in lines]
  image_ids = [x[0] for x in splitlines]
  confidence = np.array([float(x[1]) for x in splitlines])
  BB = np.array([[float(z) for z in x[2:]] for x in splitlines])
  
  conf_recs = {}
  npos_conf = 0
  for imagename in image_ids:
    R_conf = [obj for obj in recs[imagename] if obj['name'] != classname]
    class_conf = [x['name'] for x in R_conf]
    bbox_conf = np.array([x['bbox'] for x in R_conf])
    difficult_conf = np.array([x['difficult'] for x in R_conf]).astype(np.bool)
    det_conf = [False] * len(R_conf)
    npos_conf = npos_conf + sum(~difficult_conf)
    conf_recs[imagename] = {'bbox': bbox_conf,
                            'difficult': difficult_conf,
                            'det': det_conf,
                            'class' : class_conf}
  
  nd = len(image_ids)
  tp = np.zeros(nd)
  fp = np.zeros(nd)
  
  
  cls_correct = np.zeros(nd)
  cls_conf = np.zeros(nd)
  cls_name = []
  
  if BB.shape[0] > 0:
    # sort by confidence
    sorted_ind = np.argsort(-confidence)
    sorted_scores = np.sort(-confidence)
    BB = BB[sorted_ind, :]
    image_ids = [image_ids[x] for x in sorted_ind]

    # go down dets and mark TPs and FPs
    for d in range(nd):
      R = class_recs[image_ids[d]]
      bb = BB[d, :].astype(float)
      ovmax = -np.inf
      BBGT = R['bbox'].astype(float)

      if BBGT.size > 0:
        # compute overlaps
        # intersection
        ixmin = np.maximum(BBGT[:, 0], bb[0])
        iymin = np.maximum(BBGT[:, 1], bb[1])
        ixmax = np.minimum(BBGT[:, 2], bb[2])
        iymax = np.minimum(BBGT[:, 3], bb[3])
        iw = np.maximum(ixmax - ixmin + 1., 0.)
        ih = np.maximum(iymax - iymin + 1., 0.)
        inters = iw * ih

        # union
        uni = ((bb[2] - bb[0] + 1.) * (bb[3] - bb[1] + 1.) +
               (BBGT[:, 2] - BBGT[:, 0] + 1.) *
               (BBGT[:, 3] - BBGT[:, 1] + 1.) - inters)
        
        overlaps = inters / uni
        ovmax = np.max(overlaps)
        jmax = np.argmax(overlaps)
      # pdb.set_trace()
      
      if ovmax > ovthresh:
        if not R['difficult'][jmax]:
          if not R['det'][jmax]:
            cls_correct[d] = 1.
            R['det'][jmax] = 1
            cls_name.append(classname)
        else:
          cls_name.append('difficult')
      else:
        R_conf = conf_recs[image_ids[d]]
        BBGT_conf = R_conf['bbox'].astype(float)
        
        ovmax_conf = -np.inf
        if BBGT_conf.size > 0:
          # compute overlaps
          # intersection
          ixmin = np.maximum(BBGT_conf[:, 0], bb[0])
          iymin = np.maximum(BBGT_conf[:, 1], bb[1])
          ixmax = np.minimum(BBGT_conf[:, 2], bb[2])
          iymax = np.minimum(BBGT_conf[:, 3], bb[3])
          iw = np.maximum(ixmax - ixmin + 1., 0.)
          ih = np.maximum(iymax - iymin + 1., 0.)
          inters = iw * ih
  
          # union
          uni = ((bb[2] - bb[0] + 1.) * (bb[3] - bb[1] + 1.) +
                 (BBGT_conf[:, 2] - BBGT_conf[:, 0] + 1.) *
                 (BBGT_conf[:, 3] - BBGT_conf[:, 1] + 1.) - inters)
          
          overlaps_conf = inters / uni
          ovmax_conf = np.max(overlaps_conf)
          jmax_conf = np.argmax(overlaps_conf)
        
        if ovmax_conf > ovthresh:
          if not R_conf['difficult'][jmax_conf]:
            if not R_conf['det'][jmax_conf]:
              cls_conf[d] = 1.
              R_conf['det'][jmax_conf] = 1
              cls_name.append(R_conf['class'][jmax_conf])
          else:
            cls_name.append('difficult')
        else:
          cls_name.append('background')
  
  
  correct = cls_correct.mean()
  confusion = cls_conf.mean()
  
  # pdb.set_trace()
  return correct, confusion, cls_name
